#!/usr/bin/python
# -*- encoding: utf-8 -*-
import json
import os
import os.path as osp

import albumentations as A
import cv2
import numpy as np
import torchvision.transforms as transforms
from torch.utils.data import Dataset

# from transform import *


def get_img_list(path):
    img_list = []
    for root, dirs, files in os.walk(path):
        for f in files:
            if f.endswith('.jpg'):
                img_list.append(osp.join(root, f))

    assert len(img_list) > 0, "没有图片！"
    return img_list


class Steel(Dataset):
    def __init__(self, rootpth, mode='train', use_heatmap=False, inputsize=[800, 800], return_path=False, *args, **kwargs):
        super(Steel, self).__init__(*args, **kwargs)
        assert mode in ('train', 'test')
        self.mode = mode
        print('self.mode', self.mode)
        print('use_heatmap', use_heatmap)
        self.ignore_lb = 255
        self.return_path = return_path

        if use_heatmap:
            label_dir = "_heatmap"
        else:
            label_dir = "_mask"

        ## parse img directory
        img_list = get_img_list(rootpth)
        self.total_data = []
        for im_path in img_list:
            label_path = im_path.replace('.jpg', '_mask.png').replace(mode, mode + label_dir)
            assert os.path.exists(label_path), "path is not exist"
            self.total_data.append([im_path, label_path])

        self.len = len(self.total_data)
        print('mode: {} \t self.len: {}'.format(self.mode, self.len))

        ## pre-processing
        self.to_tensor = transforms.Compose([
            transforms.ToTensor(),
            # transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
            ])
        
        self.transform = A.Compose([
            A.Resize(inputsize[0], inputsize[1], interpolation=cv2.INTER_NEAREST)
        ])


    def __getitem__(self, idx):
        im_path, label_path  = self.total_data[idx]
        img = cv2.imread(im_path)
        label = cv2.imread(label_path, cv2.IMREAD_GRAYSCALE)

        im_lb = self.transform(image = img, mask = label)
        img, label = im_lb['image'], im_lb['mask']

        img = self.to_tensor(img)
        label = np.array(label / 255, np.float32)
        label = label[np.newaxis, :]
        # label = self.convert_labels(label)

        if self.return_path:
            return img, label, im_path
        else:
            return img, label

    def __len__(self):
        return self.len


    def convert_labels(self, label):
        for k, v in self.lb_map.items():
            label[label == k] = v
        return label




if __name__ == "__main__":
    from tqdm import tqdm
    ds = Steel(r'data/shushu/train', mode='train', use_heatmap=False)

    for im, lb in tqdm(ds):

        im = im.numpy().transpose(1,2,0) * 255
        im = im.astype(np.uint8)
        lb = (lb[0] * 255).astype(np.uint8)
        cv2.imshow('lb', lb)
        cv2.imshow("im", im)
        cv2.waitKey(0)

        print('ok')


